"""
-------------
captioneer.py
-------------
By Sumant Manne (darkpyro@gmail.com)
Written in Python 3.x with ten million standard imports.
GPLv3 I guess.

Warning: Does not imbue powers of quad-airshots.

"""
import codecs
import collections
import configparser
import itertools
import os
import re
import subprocess
import sys

class Captioner:
	def __init__(self, printInfo = False):
		# read config settings
		self.cfg = configparser.SafeConfigParser()
		self.cfg.readfp(open('captioneer.cfg'))
		self.locationPrefix = self.cfg.get('General', 'location-prefix')
		self.captionFiles = list(map(str.strip, self.cfg.get('General', 'files').split(',')))
		self.packageFile = os.path.join(self.locationPrefix, '../../team fortress 2 content.gcf')
		self.sdkPackageFile = os.path.join(self.locationPrefix, '../../sourcesdk.gcf')
		self.sdkCaptionCompiler = 'bin/orangebox/bin/captioncompiler.exe'
		self.sdkFileSystemSteamDll = 'bin/FileSystem_Steam.dll'
		self.outFile = os.path.join(self.locationPrefix, 'tf/resource/closecaption_english.txt')
		self.printInfo = printInfo
		self.gcfExtract = self.cfg.getboolean('General', 'copy-from-gcf')
		
		# read colors
		self.colorcfg = configparser.SafeConfigParser()
		self.colorcfg.readfp(open('captioneer-colors.cfg'))
		# actually uses modified Wolf:ET colors	
		self.quakeColors = {}
		for key, colorName in self.colorcfg.items('Quake Colors'):
			self.quakeColors[key] = colorName.strip()
		self.colors = {}
		for color, colorRGB in self.colorcfg.items('Colors'):
			self.colors[color] = tuple(colorRGB.split())
		self.print("Found {0} Qcolors and {1} colors".format(len(self.quakeColors), len(self.colors)))
		
		self.captions = list()
		self.captionAssigns = collections.OrderedDict()
		if not self.cfg.has_section('Caption Variables'):
			self.print("Warning: no Caption Variables section found")
		self.varMap = self.cfg.items('Caption Variables')

	def copyFromGCF(self):
		for captionFile in self.captionFiles:
			if not os.path.exists(os.path.join(self.locationPrefix, captionFile)) and self.gcfExtract:
				captionFileDest = os.path.dirname(os.path.join(self.locationPrefix, captionFile))
				self.HLExtractFile(self.packageFile, captionFile, captionFileDest)
	
	def HLExtractFile(self, packageFile, fileToExtract, destination):
		HLExtractCmd = 'HLExtract.exe -s -v -p "{0}" -e "{1}" -d "{2}"'.format(packageFile, fileToExtract, destination)
		try:
			retcode = subprocess.call(HLExtractCmd)
			if retcode != 0:
				self.print("failed to extract {0} from GCF".format(fileToExtract))
			else:
				self.print("Extracted {0} > {1}".format(fileToExtract, destination))
		except OSError as e:
			self.print("GCF extraction failure on {0}, ".format(packageFile) + str(e))
	
	def getCaptions(self):
		for captionFile in self.captionFiles: 
			self.getCaptionsFromFile(os.path.join(self.locationPrefix, captionFile))
	
	def getCaptionsFromFile(self, captionFile):
		with open(captionFile, 'r') as cf:
			self.print("Reading captions from {0}".format(captionFile))
			braceLevel = 0
			for line in cf:
				if line.startswith('//'):
					continue
				if line.strip().startswith('{'):
					braceLevel += 1
					continue
				if line.strip().startswith('}'):
					braceLevel -= 1
					continue
				if braceLevel == 0:
					m = re.match(r'\s*"(\S*)"\s*', line)
					if m:
						caption = m.group(1)
						self.captions.append(caption)
	
	def setCaptions(self):
		unusedCaptions = list(self.captions)
		captionRules = self.cfg.items('Caption Rules')
		self.print("Found {0} caption rules and {1} caption variables".format(len(captionRules), len(self.varMap)))
		for captionRule, captionValue in captionRules:
			for caption in self.captions:
				if re.match(captionRule + r'\d*$', caption, re.IGNORECASE):
					self.captionAssigns[caption] = captionValue
					if caption in unusedCaptions:
						unusedCaptions.remove(caption)
		if self.cfg.getboolean('General', 'default-captions'):
			for uCaption in unusedCaptions:
				self.captionAssigns[uCaption] = uCaption + '?'

	def generateCaption(self, caption):
		for var, val in self.varMap:
			caption = re.sub(r'\$\({0}\)'.format(var), val, caption)
		for qColor in self.quakeColors:
			caption = caption.replace(qColor, self.colorTag(qColor))
		return caption
	
	def writeCaptionFile(self):
		# captioncompiler demands a unicode file
		with codecs.open(self.outFile, encoding = 'utf-16', mode = 'w') as ocf:
			ocf.write('"lang"\n')
			ocf.write('{\n')
			ocf.write('\t"Language"\t"english"\n')
			ocf.write('\t"Tokens"\n')
			ocf.write('\t{\n')
			numCaptions = 0
			# write generated captions
			for ck, cv in self.captionAssigns.items():
				ocf.write('\t\t"{0}"\t\t"{1}"\n'.format(ck, self.generateCaption(cv)))
				numCaptions += 1
			ocf.write('\t\t\n')
			# write custom captions
			for cck, ccv in self.cfg.items('Custom Captions', True):
				ocf.write('\t\t"{0}"\t\t"{1}"\n'.format(cck, self.generateCaption(ccv)))
				numCaptions += 1
			ocf.write('\t}\n')
			ocf.write('}\n')
		self.print ("Wrote {0} captions to {1}".format(numCaptions, c.outFile))
	
	def callCaptionCompiler(self):
		captionCompilerPath = os.path.join(self.locationPrefix, '../sourcesdk/bin/orangebox/bin/captioncompiler.exe')
		fileSystemSteamDll = os.path.join(self.locationPrefix, '../sourcesdk/bin/orangebox/bin/FileSystem_Steam.dll')
		if not os.path.exists(captionCompilerPath):
			self.print("captioncompiler.exe not found")
			if self.gcfExtract:
				self.HLExtractFile(self.sdkPackageFile, self.sdkCaptionCompiler, os.path.dirname(captionCompilerPath))
			else:
				sys.exit(1)
		if not os.path.exists(fileSystemSteamDll):
			self.print("FileSystem_Steam.dll not found")
			if self.gcfExtract:
				self.HLExtractFile(self.sdkPackageFile, self.sdkFileSystemSteamDll, os.path.dirname(fileSystemSteamDll))
			else:
				sys.exit(1)
		self.print("Starting captioncompiler.exe...")
		try:
			retcode = subprocess.call('{0} "{1}"'.format(captionCompilerPath, os.path.abspath(self.outFile)))
			if retcode < 0:
				self.print("Caption compiling failed!")
				sys.exit(retcode)
			else:
				self.print("Captions creation succeeded!")
		except OSError as e:
			self.print("Caption compiling error, " + str(e))
			sys.exit(1)
		
	def colorTag(self, color):
		if color in self.quakeColors:
			color = self.quakeColors[color]
		if color in self.colors:
			r, g, b = self.colors[color]
			return '<clr:{0},{1},{2}>'.format(r, g, b)
		else:
			return ''
	
	def print(self, str):
		if self.printInfo:
			print(str)

if __name__ == '__main__':
	c = Captioner(True)
	c.copyFromGCF()
	c.getCaptions()
	c.setCaptions()
	c.writeCaptionFile()
	c.callCaptionCompiler()
